"""
    启动服务器时，使用以下命令：
    uvicorn main:app --reload
    require: fastapi, pymilvus, face_recognize, opencv-python, numpy, uvicorn, websockets
"""

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.websockets import WebSocketState
from pymilvus import MilvusClient
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import face_recognition
import uvicorn
import cv2
import asyncio
import base64
import numpy as np

# 定义请求体数据模型
class ImageRequest(BaseModel):
    userId: int
    file: str

# 连接 Milvus 数据库
client = MilvusClient(
    uri="http://localhost:19530",
    db_name="face_db"
)
print("连接MilvusClient成功")
# 加载集合
client.load_collection(
    collection_name="face_collection"
)
print("加载集合成功")

app = FastAPI()
connected_clients = []

# WebSocket 接口：用于接收图像并进行人脸识别
@app.websocket("/ws/recognize")
async def recognize(websocket: WebSocket):
    """
    人脸识别
    """
    await websocket.accept()
    connected_clients.append(websocket)
    try:
        asyncio.create_task(close_connection_after_delay(websocket, 10))

        while True:
            data = await websocket.receive_text()  # 接收前端发送的图像数据
            data = base64.b64decode(data)
            np_array = np.frombuffer(data, dtype=np.uint8)  # 转换为 NumPy 数组
            image = cv2.imdecode(np_array, cv2.IMREAD_COLOR)  # 解码为图像

            # 保存图片
            # cv2.imwrite("/home/wei/Desktop/WS/wx.jpg", image)

            if image is None:
                raise HTTPException(status_code=400, detail="Invalid image data.")

            # 使用 face_recognition 检测人脸
            face_locations = face_recognition.face_locations(image)
            if len(face_locations) == 0:
                await websocket.send_json({"code": 400, "message": "No faces detected!", "data": None})  # 返回识别结果
            elif len(face_locations) > 1:
                await websocket.send_json({"code": 400, "message": "Too much faces detected!", "data": None})  # 返回识别结果
            else:
                # 提取特征向量
                face_encodings = face_recognition.face_encodings(image, face_locations)

                query_vector = face_encodings[0]

                res = client.search(
                    collection_name="face_collection",
                    anns_field="face_info",
                    data=[query_vector],
                    limit=3,
                    output_fields = ["face_info"]
                )

                user_id = -1

                for hits in res:
                    for hit in hits:
                        face_info = hit['entity']['face_info']
                        if face_recognition.compare_faces(np.array(face_encodings), face_info, 0.5)[0]:
                           user_id = hit['id']
                           break

                if user_id == -1:
                    await websocket.send_json({"code": 401, "message": "No faces message registed!", "data": None})  # 返回识别结果
                else:
                    await websocket.send_json({"code": 200, "message": "success", "data": user_id})

    except (WebSocketDisconnect, HTTPException, RuntimeError):
        print("Client disconnected.")
    finally:
        connected_clients.remove(websocket)
        print('成功移除websocket')

@app.post("/register")
async def register(request_body: ImageRequest):
    """
    接收图像文件进行人脸注册
    """
    user_id = request_body.userId
    file = request_body.file
    print(user_id)
    file = base64.b64decode(file)
    try:
        # 将图像文件转换为 NumPy 数组
        np_array = np.frombuffer(file, dtype=np.uint8)
        frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)  # 解码为图像

        # 保存图片
        # cv2.imwrite("/home/wei/Desktop/WS/wx.jpg", frame)

        if frame is None:
            raise HTTPException(status_code=400, detail="Invalid image data.")

        # 使用 face_recognition 检测人脸
        face_locations = face_recognition.face_locations(frame)

        if len(face_locations) == 0:
            return JSONResponse({"code": 400, "message": "No faces detected!", "data": None})
        elif len(face_locations) > 1:
            return JSONResponse({"code": 400, "message": "Too much faces detected!", "data": None})

        # 提取特征向量
        face_encodings = face_recognition.face_encodings(frame, face_locations)

        data = []

        data.append({"user_id": user_id, "face_info": face_encodings[0]})

        client.insert(
            collection_name="face_collection",
            data=data
        )

        return JSONResponse({"code": 200, "message": "succeess", "data": None})

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

async def close_connection_after_delay(websocket: WebSocket, delay: int):
    await asyncio.sleep(delay)
    if websocket.client_state != WebSocketState.DISCONNECTED:
        await websocket.close()  # 10秒后主动关闭 WebSocket 连接
        print("Closing connection after 10 seconds")

if __name__ == "__main__":
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)
